//
// Created by liugang on 9/12/25.
//

#ifndef NODES_H
#define NODES_H

#include "list_head.h"

#include <cstdint>
#include <string>
#include <variant>

#define NODEKIND_TABLE()\
    DEF_NODEKIND(10, NK_EXPR_LITERAL, "") \
    DEF_NODEKIND(11, NK_EXPR_VAR, "求导变量") \
    DEF_NODEKIND(12, NK_EXPR_REF, "具体量的引用") \
    /*DEF_NODEKIND(13, NK_EXPR_VAL, "具体量的引用")*/ \
    \
    DEF_NODEKIND(20, NK_EXPR_ADD, "") \
    DEF_NODEKIND(21, NK_EXPR_SUB, "") \
    DEF_NODEKIND(22, NK_EXPR_MULT, "") \
    DEF_NODEKIND(23, NK_EXPR_DIV, "") \
    DEF_NODEKIND(25, NK_EXPR_SIN, "") \
    DEF_NODEKIND(26, NK_EXPR_COS, "") \
    /* DEF_NODEKIND(29, NK_EXPR_ACCESS, "") */ \
    DEF_NODEKIND(30, NK_EXPR_GT, "") \
    DEF_NODEKIND(31, NK_EXPR_LT, "") \
    DEF_NODEKIND(32, NK_EXPR_EQ, "") \
    DEF_NODEKIND(33, NK_EXPR_PHI, "") \
    DEF_NODEKIND(34, NK_EXPR_ASSIGN, "") \
    \
    DEF_NODEKIND(50, NK_STMT_STOP, "Stop execution") \
    DEF_NODEKIND(51, NK_STMT_END, "Block Ended") \
    DEF_NODEKIND(52, NK_STMT_JUMP, "Jump to specified node") \
    DEF_NODEKIND(53, NK_STMT_BLOCK, "") \
    DEF_NODEKIND(54, NK_STMT_EXPR, "") \
    /* (end) */

using real = double;
using integer = double;

enum NodeKind {
#define DEF_NODEKIND(Num_, Name_, Desc_) Name_ = Num_,
    NODEKIND_TABLE()
#undef DEF_NODEKIND
};

enum NodeFlags : uint32_t {
    // clang-format off
    NF_EXPR                 = 0x0001, // 属于表达式
    NF_STMT                 = 0x0002, // 属于语句
    NF_VAR                  = 0x0004, // 属于求导变量

    // 块的类型
    NF_BLOCK_MASK           = 0x0070,   // 块类型掩码
    NF_BLOCK_SEQ            = 0x0000,   // 普通序列块
    NF_BLOCK_IF             = 0x0010,   // IF 块
    NF_BLOCK_FOR            = 0x0020,   // FOR 块

    // Jump 的种类
    NF_JUMP_MASK            = 0x0700,   // JUMP 类型掩码
    NF_JUMP_FORCE           = 0x0000,   // 不带提交的跳转
    NF_JUMP_IF_COND         = 0x0010,   // if 的条件表达式中的跳转
    NF_JUMP_FOR_COND        = 0x0200,   // for 的条件表达式中的跳转
    NF_JUMP_FOR_BREAK       = 0x0300,   // for 中的 break 导致的跳转
    NF_JUMP_FOR_CONTINUE    = 0x0400,   // for 中的 continue 导致的跳转

    // 节点的性质
    NF_NONDIFFERENTIABLE    = 0x0800,   // 不可求导的
    // clang-format on
};

// FOR
// +-- init_stmt
// +-- JUMP:END   <-----+
//     +-- cond_expr    |
// +-- body_stmt        |
// +-- iter_stmt        |
// +-- JUMP[continue]---+
// +-- END

// IF
// +-- JUMP:END
//     +-- cond_expr
// +-- true_stmt
// +-- JUMP:END
// +-- false_stmt
// +-- END

struct Node {
    Node* next{};
    Node* prev{};
    Node* parent{};
    list_head children{};

    NodeKind kind{};
    uint32_t flags{0};

    real cached_value{};

    explicit Node(NodeKind k, uint32_t f = 0u)
        : kind(k), flags(f) {
        list_head_init(this);
    }

    [[nodiscard]] auto is_expr() const -> bool {
        return !!(flags & NF_EXPR);
    }

    [[nodiscard]] auto is_stmt() const -> bool {
        return !!(flags & NF_STMT);
    }

    [[nodiscard]] auto is_nondifferentiable() const -> bool {
        return !!(flags & NF_NONDIFFERENTIABLE);
    }

    [[nodiscard]] auto is_differentiable() const -> bool {
        return !is_nondifferentiable();
    }

    auto add_child(Node* child) -> void {
        ASSERT(child!= nullptr);

        list_append(&children, reinterpret_cast<list_head*>(child));
        child->parent = this;

        // 不可求导是可以被继承的
        if (is_nondifferentiable()) {
            child->flags |= NF_NONDIFFERENTIABLE;
        }
    }

    auto insert_child(Node* prev, Node* next, Node* child) -> void {
        ASSERT(child!= nullptr);

        ASSERT(prev != nullptr);
        ASSERT(next != nullptr);
        ASSERT(prev->parent == this);
        ASSERT(next->parent == this);

        list_insert(prev, next, child);

        child->parent = this;

        // 不可求导是可以被继承的
        if (is_nondifferentiable()) {
            child->flags |= NF_NONDIFFERENTIABLE;
        }
    }

    [[nodiscard]] auto child(int32_t index) -> Node* {
        return reinterpret_cast<Node*>(list_child(&children, index));
    }

    [[nodiscard]] auto child(int32_t index) const -> const Node* {
        return reinterpret_cast<const Node*>(list_child(&children, index));
    }

    [[nodiscard]] auto empty() const -> bool {
        return list_empty(&children);
    }
};

struct Expr : public Node {
    explicit Expr(NodeKind k, uint32_t f = 0u)
        : Node{k, f} {
        flags |= NF_EXPR;
    }
};

struct Stmt : public Node {
    explicit Stmt(NodeKind k, uint32_t f = 0u, Node* child = nullptr)
        : Node{k, f} {
        flags |= NF_STMT;
        if (child != nullptr) {
            add_child(child);
        }
    }
};

struct StmtJump : public Stmt {
    Node* target{};

    explicit StmtJump(uint32_t f, Node* t, Node* cond_expr = nullptr)
        : Stmt(NK_STMT_JUMP, f, cond_expr) {
        target = t;
    }
};

struct ExprLiteral : public Expr {
    real val{};

    explicit ExprLiteral(real v)
        : Expr{NK_EXPR_LITERAL, NF_NONDIFFERENTIABLE}, val{v} {
    }
};


class FLOW;

struct ExprVar : public Expr {
    FLOW* flow{};
    int32_t access_proc_id;
    int32_t loc_id{-1};

    explicit ExprVar(FLOW* f, int32_t pid, int32_t lid)
        : Expr{NK_EXPR_VAR}, flow{f}, access_proc_id(pid), loc_id(lid) {
    }
};

struct ExprRef : public Expr {
    real* addr{};

    explicit ExprRef(real& v)
        : Expr{NK_EXPR_REF, NF_NONDIFFERENTIABLE}, addr{&v} {
    }
};

// struct ExprVal : public Expr {
//     real* addr{};
//
//     explicit ExprVal(real& v)
//         : Expr{NK_EXPR_VAL}, addr{&v} {
//     }
// };


// struct ExprAccess : public Expr {
//     FLOW* flow;
//     int32_t access_proc_index;
//     int32_t local_id;
//
//     explicit ExprAccess(FLOW* f, int32_t idx, int32_t loc_id)
//         : Expr{NK_EXPR_ACCESS}, flow{f}, access_proc_index{idx}, local_id(loc_id) {
//     }
// };

enum PrintFlags {
    // 是否打印子节点
    PF_PRINTCHILDREN = 0x1,
};

extern auto print_node(const Node* node, int level, uint16_t flags, FILE* file) -> void;

extern auto eval_expr(FLOW* flow, Node* n) -> real;
extern auto eval_stmt(FLOW* flow, Node* n) -> void;

extern auto deriv_expr(FLOW* flow, Node* n, real derval) -> void;
extern auto deriv_stmt(FLOW* flow, Node* n, real derval) -> void;
#endif //NODES_H